

#include "amix.h"
#include "base_util.h"
#include "packet_frame_operate.h"
#include <cstdio>
#include <libavutil/frame.h>
#include <thread>
#include <vector>

int WriteAudioData(AVFrame *frame) {

    // 多少位的编码
    int data_size = av_get_bytes_per_sample((AVSampleFormat)frame->format);
    if (data_size < 0) {

        return -1;
    }
    FILE * file_handle = fopen("30.pcm", "a+");
    if (!file_handle)
    {
        return 1;
    }
    // frame->data 中的音频数据是以packet形式放置的,在原始pcm数据中是以planer放置的
    const int n = frame->nb_samples * av_get_channel_layout_nb_channels(frame->channel_layout);
    const uint16_t *p     = (uint16_t*)frame->data[0];
    const uint16_t *p_end = p + n;

    while (p < p_end) {
        fputc(*p    & 0xff, file_handle);
        fputc(*p>>8 & 0xff, file_handle);
        p++;
    }
    fclose(file_handle);
    return 0;
}

void InitReadPacketDataThread(AudioMix & mix_data)
{
    for (auto & data : mix_data.m_input_file_map)
    {
        auto input_file = data.second;
        input_file->m_read_packet_queue.InitMessageQueue(1000);
        input_file->m_decode_frame_queue.InitMessageQueue(1000);
        input_file->m_reach_end = false;
        input_file->m_thread = std::thread([input_file](){

            AVPacket * packet = av_packet_alloc();
            ReadPacketDataToMessageQueue(&input_file->m_input_ctx, packet, input_file->m_read_packet_queue);
            av_packet_free(&packet);
        });
    }
}

void MainLoop(AudioMix & mix_data)
{
    int select_stream = -1;
    while (true)
    {
        select_stream = -1;
        int64_t max_time = INT64_MAX;
        for (auto & data : mix_data.m_input_file_map)
        {
            if (!data.second->m_reach_end && data.second->m_input_time < max_time)
            {
                select_stream = data.first;
                max_time      = data.second->m_input_time;
            }
        }
        if (select_stream == -1)
        {
            break;
        }
        auto select_file = mix_data.m_input_file_map[select_stream];
        AVPacket * packet = nullptr;
        int result = select_file->m_read_packet_queue.PopPacket(&packet);
        if (result == AVERROR_EOF)
        {
            if (packet)
            {
                DecodePacketDataToMessageQueue(select_file->m_input_ctx.AudioDecoder(), packet, select_file->m_decode_frame_queue);
            }
            DecodePacketDataToMessageQueue(select_file->m_input_ctx.AudioDecoder(), nullptr, select_file->m_decode_frame_queue);
            select_file->m_reach_end = true;
        }
        else if (packet)
        {
            DecodePacketDataToMessageQueue(select_file->m_input_ctx.AudioDecoder(), packet, select_file->m_decode_frame_queue);
            select_file->m_input_time += packet->duration * av_q2d(select_file->m_input_ctx.AudioDecoder()->AvCodecCtx()->pkt_timebase) * AV_TIME_BASE;
        }
        // send frame data to audio filter
        ProcessFrameDataFromMessageQueue(select_file->m_decode_frame_queue, [&mix_data,&select_file](AVFrame * frame){

            FilterProcessFrame(mix_data, select_file, frame);

        }, false);
        if (select_file->m_reach_end)
        {
            FilterProcessFrame(mix_data, select_file, nullptr);
        }
        ProcessFrameDataFromMessageQueue(mix_data.m_filter_frame_queue, [&mix_data](AVFrame * frame){

            //printf("pts : %ld\tformat : %d\tchannel : %d from file : %d.\n", frame->pts, frame->format, frame->channels, select_file->m_index);
            if (!mix_data.m_output_ctx.AvFormatContext())
            {
                InitOutputFile(mix_data, frame);
            }
            // encode audio data
            EncodeFrameDataToMessageQueue(mix_data.m_output_ctx.AudioEncoder(), frame, mix_data.m_encode_packet_queue);
            ProcessPacketDataFromMessageQueue(mix_data.m_encode_packet_queue, [&mix_data](AVPacket * packet){

                auto filter_tm = av_buffersink_get_time_base(mix_data.m_buffer_sink_ctx);
                av_packet_rescale_ts(packet, filter_tm, mix_data.m_output_ctx.AudioMediaStream()->AvStream()->time_base);
                packet->time_base = mix_data.m_output_ctx.AudioMediaStream()->AvStream()->time_base;
                packet->stream_index = mix_data.m_output_ctx.AudioMediaStream()->AvStream()->index;
                av_interleaved_write_frame(mix_data.m_output_ctx.AvFormatContext(), packet);

            }, false);
        }, false);
        if (packet)
        {
            av_packet_unref(packet);
            av_packet_free(&packet);
        }
    }
    while (true)
    {
        AVFrame * frame = nullptr;
        mix_data.m_filter_frame_queue.PopFrameNonBlocking(&frame);
        if (!frame)
        {
            break;
        }
        EncodeFrameDataToMessageQueue(mix_data.m_output_ctx.AudioEncoder(), frame, mix_data.m_encode_packet_queue);
    }
    EncodeFrameDataToMessageQueue(mix_data.m_output_ctx.AudioEncoder(), nullptr, mix_data.m_encode_packet_queue);
    ProcessPacketDataFromMessageQueue(mix_data.m_encode_packet_queue, [&mix_data](AVPacket * packet){

        auto filter_tm = av_buffersink_get_time_base(mix_data.m_buffer_sink_ctx);
        av_packet_rescale_ts(packet, filter_tm, mix_data.m_output_ctx.AudioMediaStream()->AvStream()->time_base);
        packet->time_base = mix_data.m_output_ctx.AudioMediaStream()->AvStream()->time_base;
        packet->stream_index = mix_data.m_output_ctx.AudioMediaStream()->Index();
        av_interleaved_write_frame(mix_data.m_output_ctx.AvFormatContext(), packet);

    }, false);
    av_write_trailer(mix_data.m_output_ctx.AvFormatContext());
}

int main(int argc, const char * argv[])
{
    // av_log_set_level(AV_LOG_VERBOSE);
    AudioMix audio_mix;
    std::vector<std::string> input_file_vec = {"/home/hjie/source_project/open_net/ffmpeg/resource/sample3.aac", "/home/hjie/source_project/open_net/ffmpeg/resource/sample3.opus"};
    if (!InitInputFile(input_file_vec,audio_mix))
    {
        return -1;
    }
    // create
    InitReadPacketDataThread(audio_mix);
    // main loop
    MainLoop(audio_mix);
    for (auto & data : audio_mix.m_input_file_map)
    {
        if (data.second->m_thread.joinable())
        {
            data.second->m_thread.join();
        }
    }
    return EXIT_SUCCESS;
}
